#!/bin/env python

import time
import cv2
import numpy as np

# 图像的宽度和高度
width = 2560
height = 1440
nv12_file_path = "2k.nv12"

width = 1920
height = 1080
nv12_file_path = "1080p.nv12"

width = 640
height = 360
nv12_file_path = "360p.nv12"

# 读取 NV12 文件内容
with open(nv12_file_path, 'rb') as f:
    nv12_data = f.read()

# 将 NV12 数据转换为 NumPy 数组
yuv = np.frombuffer(nv12_data, np.uint8).reshape(height * 3 // 2, width)

# 分离 Y 和 UV 平面
y = yuv[0:height].reshape(height, width, 1)
uv = yuv[height:].reshape(height // 2, width // 2, 2)

# 将 UV 数据的大小调整为与 Y 数据相同
uv = cv2.resize(uv, (width, height), interpolation=cv2.INTER_NEAREST)

# 将 YUV 数据合并为一个数组
yuv_merge = np.concatenate((y, uv), axis=2)

# print("shape:", yuv_merge.shape)
# print("size:", yuv_merge.size)
# print("dtype:", yuv_merge.dtype)

# 将 YUV 数据转换为 BGR 格式
bgr = cv2.cvtColor(yuv_merge, cv2.COLOR_YUV2BGR)

print("shape:", bgr.shape)
print("size:", bgr.size)
print("dtype:", bgr.dtype)

windowName = 'YUV IMG'
cv2.namedWindow(windowName, cv2.WINDOW_NORMAL)
cv2.resizeWindow(windowName, width, height)
cv2.imshow(windowName, bgr)

windowName = 'Y IMG'
cv2.namedWindow(windowName, cv2.WINDOW_NORMAL)
cv2.resizeWindow(windowName, width, height)
cv2.imshow(windowName, y)

while True:
    if cv2.waitKey(0) == ord('q'):
        break

print('销毁窗口')
cv2.destroyAllWindows()